
from lib.storage.factory import SqliteFactory, Connect, NodeInfo


class Storage(object):

    def __init__(self):
        self.storage = SqliteFactory()

    def list(self):
        conns = []
        session = self.storage.session()
        connects = session.query(Connect).all()
        for conn in connects:
            conn_dict = dict()
            conn_dict['id'] = conn.id
            conn_dict['name'] = conn.name
            conn_dict['type'] = conn.type
            node_info_list = session.query(NodeInfo).filter(NodeInfo.connect_id == conn.id).all()
            if conn.type == 0:
                node_info = node_info_list[0]
                conn_dict['node'] = (node_info.host, node_info.port)
            elif conn.type == 1:
                nodes = []
                for node_info in node_info_list:
                    nodes.append((node_info.host, node_info.port))
                conn_dict['nodes'] = nodes
            conn_dict['auth'] = conn.auth
            conns.append(conn_dict)
        return conns

    def get(self, id):
        session = self.storage.session()
        connect = session.query(Connect).filter(Connect.id == id).first()
        if not connect:
            return None
        node_info_list = session.query(NodeInfo).filter(NodeInfo.connect_id == connect.id).all()
        conn_dict = {'id': connect.id, 'name': connect.name, 'auth': connect.auth, 'type': connect.type}
        if connect.type == 0:
            node_info = node_info_list[0]
            conn_dict['node'] = (node_info.host, node_info.port)
        elif connect.type == 1:
            nodes = []
            for node_info in node_info_list:
                nodes.append((node_info.host, node_info.port))
            conn_dict['nodes'] = nodes
        else:
            return None
        return conn_dict

    def insert(self, name, type, nodes, auth):
        session = self.storage.session()
        try:
            connect = Connect()
            connect.name = name
            connect.type = type
            connect.auth = auth
            session.add(connect)
            session.flush()
            print('connect id=', connect.id)
            for node in nodes:
                host, port = node
                node_info = NodeInfo()
                node_info.connect_id = connect.id
                node_info.host = host
                node_info.port = port
                session.add(node_info)
            session.commit()
        except Exception as e:
            print(e)
            session.rollback()
        finally:
            session.close()

    def delete(self, id):
        session = self.storage.session()
        try:
            session.query(Connect).filter(Connect.id == id).delete()
            session.query(NodeInfo).filter(NodeInfo.connect_id == id).delete()
            session.commit()
        except Exception as e:
            print(e)
            session.rollback()
        finally:
            session.close()

